# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/27 11:53 
@Author  : ZhangShenao 
@File    : 3.GPT-4o接入DALL·E文生图.py 
@Desc    : GPT-4o接入DALL·E文生图

执行流程:
1. 将用户的绘图Prompt传递给LLM,并强制调用openai_dalle绘图工具
2. LLM推理出需要调用工具,生成dalle的绘图Prompt
3. 用户将LLM生成的绘图Prompt传递给openai_dalle工具,进行实际调用,并获取图片
"""
import os

import dotenv
from langchain_openai import ChatOpenAI

from dalle_image_tool import create_dalle_image_tool

# 加载环境变量
dotenv.load_dotenv()

# 创建DALLE文生图工具
dalle_tool = create_dalle_image_tool()

# 创建LLM,并绑定工具
llm = ChatOpenAI(model="gpt-4o", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
llm_with_tools = llm.bind_tools(tools=[dalle_tool], tool_choice="openai_dalle")  # 强制LLM调用openai_dalle工具

# 构造Chain
chain = llm_with_tools | (lambda resp: resp.tool_calls[0].get("args")) | dalle_tool

# 执行Chain,获取结果
result = chain.invoke("请帮我画一张图片，一只苗条的英短蓝猫正趴在窗台上，一边望着窗外树上的小鸟，一边喵喵叫个不停")
print(result)
